{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Udb-wL0GEJ5j"
      },
      "source": [
        "##### Copyright 2019 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "5oYucXhdRXlW"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "3GzW9UkQRg7L"
      },
      "source": [
        "# Light interaction with materials\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/graphics/blob/master/tensorflow_graphics/notebooks/reflectance.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/graphics/blob/master/tensorflow_graphics/notebooks/reflectance.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "N78wnH25iV5m"
      },
      "source": [
        "The world around us is very complex and is made of a wide array of materials ranging from glass to wood. Each material possesses its own intrinsic properties and interacts differently with light. For instance, some are diffuse (e.g. paper or marble) and given a lighting condition, look the same from any angle. Other materials (e.g. metal) have an appearance that can vary significantly and exhibit view dependent effects such as specularities.\n",
        "\n",
        "Modelling exactly how light interacts with materials is a complex process that involves effects like sub-surface scattering (e.g. skin) and refraction (e.g. water). In this Colab, we focus on the most common effect which is reflection. [Bidirectional reflectance distribution functions (BRDF)](https://en.wikipedia.org/wiki/Bidirectional_reflectance_distribution_function) is the method of choice when it comes to modelling reflectance. Given the direction of incoming light, BRDFs control the amount of light that bounces in the direction the surface is being observed (any gray vector in the image below).\n",
        "\n",
        "![](https://storage.googleapis.com/tensorflow-graphics/notebooks/reflectance/brdf.jpg)\n",
        "\n",
        "In this Colab, a light we be shone onto three spheres, each with a material described in the image above, where the specular material is going to be modelled with the [Phong specular model](https://en.wikipedia.org/wiki/Phong_reflection_model).\n",
        "\n",
        "**Note**: This Colab covers an advanced topic and hence focuses on providing a controllable toy example to form a high level understanding of BRDFs rather than providing step by step details. For those interested, these details are nevertheless available in the code."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "E87e3hO5SWN9"
      },
      "source": [
        "## Setup \u0026 Imports\n",
        "If Tensorflow Graphics is not installed on your system, the following cell can install the Tensorflow Graphics package for you."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "t60P7UZASZTy"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow_graphics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Fx_oEf1ySc-V"
      },
      "source": [
        "Now that Tensorflow Graphics is installed, let's import everything needed to run the demo contained in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "D8jOq5AMcB8b"
      },
      "outputs": [],
      "source": [
        "import math as m\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow_graphics.rendering.reflectance import lambertian\n",
        "from tensorflow_graphics.rendering.reflectance import phong\n",
        "from tensorflow_graphics.rendering.camera import orthographic\n",
        "from tensorflow_graphics.geometry.representation import grid\n",
        "from tensorflow_graphics.geometry.representation import ray\n",
        "from tensorflow_graphics.geometry.representation import vector"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n5SYoy4pXRms"
      },
      "source": [
        "## Controllable lighting of a sphere"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "fAbPR3iyv_z7"
      },
      "outputs": [],
      "source": [
        "###############\n",
        "# UI controls #\n",
        "###############\n",
        "#@title Controls { vertical-output: false, run: \"auto\" }\n",
        "light_x_position = -0.4  #@param { type: \"slider\", min: -1, max: 1 , step: 0.05 }\n",
        "albedo_red = 0.7  #@param { type: \"slider\", min: 0.0, max: 1.0 , step: 0.1 }\n",
        "albedo_green = 1  #@param { type: \"slider\", min: 0.0, max: 1.0 , step: 0.1 }\n",
        "albedo_blue = 1  #@param { type: \"slider\", min: 0.0, max: 1.0 , step: 0.1 }\n",
        "light_red = 1  #@param { type: \"slider\", min: 0.0, max: 1.0 , step: 0.1 }\n",
        "light_green = 1  #@param { type: \"slider\", min: 0.0, max: 1.0 , step: 0.1 }\n",
        "light_blue = 1  #@param { type: \"slider\", min: 0.0, max: 1.0 , step: 0.1 }\n",
        "specular_percentage = 0.25  #@param { type: \"slider\", min: 0, max: 1 , step: 0.01 }\n",
        "shininess = 4  #@param { type: \"slider\", min: 0, max: 10, step: 1 }\n",
        "diffuse_percentage = 1.0 - specular_percentage\n",
        "dtype = np.float64\n",
        "albedo = np.array((albedo_red, albedo_green, albedo_blue), dtype=dtype)\n",
        "\n",
        "def compute_intersection_normal_sphere(image_width, image_height, sphere_radius,\n",
        "                                       sphere_center, dtype):\n",
        "  pixel_grid_start = np.array((0.5, 0.5), dtype=dtype)\n",
        "  pixel_grid_end = np.array((image_width - 0.5, image_height - 0.5), dtype=dtype)\n",
        "  pixel_nb = np.array((image_width, image_height))\n",
        "  pixels = grid.generate(pixel_grid_start, pixel_grid_end, pixel_nb)\n",
        "\n",
        "  pixel_ray = tf.math.l2_normalize(orthographic.ray(pixels), axis=-1)\n",
        "  zero_depth = np.zeros([image_width, image_height, 1])\n",
        "  pixels_3d = orthographic.unproject(pixels, zero_depth)\n",
        "\n",
        "  intersections_points, normals = ray.intersection_ray_sphere(\n",
        "      sphere_center, sphere_radius, pixel_ray, pixels_3d)\n",
        "  intersections_points = np.nan_to_num(intersections_points)\n",
        "  normals = np.nan_to_num(normals)\n",
        "  return intersections_points[0, :, :, :], normals[0, :, :, :]\n",
        "\n",
        "#####################################\n",
        "# Setup the image, sphere and light #\n",
        "#####################################\n",
        "# Image dimensions\n",
        "image_width = 400\n",
        "image_height = 300\n",
        "\n",
        "# Sphere center and radius\n",
        "sphere_radius = np.array((100.0,), dtype=dtype)\n",
        "sphere_center = np.array((image_width / 2.0, image_height / 2.0, 300.0),\n",
        "                         dtype=dtype)\n",
        "\n",
        "# Set the light along the image plane\n",
        "light_position = np.array((image_width / 2.0 + light_x_position * image_width,\n",
        "                           image_height / 2.0, 0.0),\n",
        "                          dtype=dtype)\n",
        "vector_light_to_sphere_center = light_position - sphere_center\n",
        "light_intensity_scale = vector.dot(\n",
        "    vector_light_to_sphere_center, vector_light_to_sphere_center,\n",
        "    axis=-1) * 4.0 * m.pi\n",
        "light_intensity = np.array(\n",
        "    (light_red, light_green, light_blue)) * light_intensity_scale\n",
        "\n",
        "################################################################################################\n",
        "# For each pixel in the image, estimate the corresponding surface point and associated normal. #\n",
        "################################################################################################\n",
        "intersection_3d, surface_normal = compute_intersection_normal_sphere(\n",
        "    image_width, image_height, sphere_radius, sphere_center, dtype)\n",
        "\n",
        "#######################################\n",
        "# Reflectance and radiance estimation #\n",
        "#######################################\n",
        "incoming_light_direction = tf.math.l2_normalize(\n",
        "    intersection_3d - light_position, axis=-1)\n",
        "outgoing_ray = np.array((0.0, 0.0, -1.0), dtype=dtype)\n",
        "albedo = tf.broadcast_to(albedo, tf.shape(surface_normal))\n",
        "\n",
        "# Lambertian BRDF\n",
        "brdf_lambertian = diffuse_percentage * lambertian.brdf(incoming_light_direction, outgoing_ray,\n",
        "                                  surface_normal, albedo)\n",
        "# Phong BRDF\n",
        "brdf_phong = specular_percentage * phong.brdf(incoming_light_direction, outgoing_ray, surface_normal,\n",
        "                        np.array((shininess,), dtype=dtype), albedo)\n",
        "# Composite BRDF\n",
        "brdf_composite = brdf_lambertian + brdf_phong\n",
        "# Irradiance\n",
        "cosine_term = vector.dot(surface_normal, -incoming_light_direction)\n",
        "cosine_term = tf.math.maximum(tf.zeros_like(cosine_term), cosine_term)\n",
        "vector_light_to_surface = intersection_3d - light_position\n",
        "light_to_surface_distance_squared = vector.dot(\n",
        "    vector_light_to_surface, vector_light_to_surface, axis=-1)\n",
        "irradiance = light_intensity / (4 * m.pi *\n",
        "                                light_to_surface_distance_squared) * cosine_term\n",
        "# Rendering equation\n",
        "zeros = tf.zeros(intersection_3d.shape)\n",
        "radiance = brdf_composite * irradiance\n",
        "radiance_lambertian = brdf_lambertian * irradiance\n",
        "radiance_phong = brdf_phong * irradiance\n",
        "\n",
        "###############################\n",
        "# Display the rendered sphere #\n",
        "###############################\n",
        "# Saturates radiances at 1 for rendering purposes.\n",
        "radiance = np.minimum(radiance, 1.0)\n",
        "radiance_lambertian = np.minimum(radiance_lambertian, 1.0)\n",
        "radiance_phong = np.minimum(radiance_phong, 1.0)\n",
        "# Gamma correction\n",
        "radiance = np.power(radiance, 1.0 / 2.2)\n",
        "radiance_lambertian = np.power(radiance_lambertian, 1.0 / 2.2)\n",
        "radiance_phong = np.power(radiance_phong, 1.0 / 2.2)\n",
        "\n",
        "plt.figure(figsize=(20, 20))\n",
        "\n",
        "# Diffuse\n",
        "radiance_lambertian = np.transpose(radiance_lambertian, (1, 0, 2))\n",
        "ax = plt.subplot(\"131\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Lambertian\")\n",
        "_ = ax.imshow(radiance_lambertian)\n",
        "\n",
        "# Specular\n",
        "radiance_phong = np.transpose(radiance_phong, (1, 0, 2))\n",
        "ax = plt.subplot(\"132\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Specular - Phong\")\n",
        "_ = ax.imshow(radiance_phong)\n",
        "\n",
        "# Diffuse + specular\n",
        "radiance = np.transpose(radiance, (1, 0, 2))\n",
        "ax = plt.subplot(\"133\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Combined lambertian and specular\")\n",
        "_ = ax.imshow(radiance)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "reflectance.ipynb",
      "provenance": [],
      "toc_visible": true,
      "private_outputs": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
